from tools_algorithm import *


class Node:
    def __init__(self) -> None:
        self.h = {}
        self.cnt = 0


class Solution:
    def countPrefixSuffixPairs(self, words: List[str]) -> int:
        ans = 0
        root = Node()

        for word in words:
            n = len(word)
            z = [0] * n
            l, r = 0, 0
            rec = set([n - 1])
            for i in range(1, n):
                z[i] = 0 if i > r else min(z[i - l], r - i + 1)
                if z[i] == n - i:
                    rec.add(n - i - 1)
                while (j := i + z[i]) < n and word[z[i]] == word[j]:
                    l, r = i, j
                    z[i] += 1
                    if z[i] == n - i:
                        rec.add(n - i - 1)
            node = root
            for i, v in enumerate(word):
                if v not in node.h:
                    break
                node = node.h[v]
                if i in rec:
                    ans += node.cnt
            node = root
            for v in word:
                if v not in node.h:
                    node.h[v] = Node()
                node = node.h[v]
            node.cnt += 1
        return ans


s = Solution()
print(s.countPrefixSuffixPairs(["a", "aba", "ababa", "aa"]))
